import numpy as np
import torch
import random


class ToTensor:
    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        print(img)
        img = np.array(img)
        print(img)
        img = img.astype(np.float32).transpose((2, 0, 1))
        mask = np.array(mask)
        mask = mask.astype(np.float32)

        img = torch.from_numpy(img).float()
        mask = torch.from_numpy(mask).float()

        return {'image': img, 'label': mask}


class Normalize:
    def __init__(self, mean=(0., 0., 0.), std=(1., 1., 1.)):
        self.mean = mean
        self.std = std

    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        img = np.array(img).astype(np.float32)
        mask = np.array(mask).astype(np.float32)
        img /= 255.0
        img -= self.mean
        img /= self.std

        return {'image': img, 'label': mask}


class RandomCrop:
    def __init__(self, width, height):
        self.width = width
        self.height = height

    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']

        w, h = img.size

        x_scale = w - self.width
        y_scale = h - self.height

        x = random.randint(0, x_scale)
        y = random.randint(0, y_scale)

        img = img.crop((x, y, x + self.width, y + self.height))
        mask = mask.crop((x, y, x + self.width, y + self.height))
        return {'image': img, 'label': mask}


